#ifndef COARSEN_HEAP_HPP
#define COARSEN_HEAP_HPP
#include "utils.hpp"
#include "coarsen.hpp"
#include <unordered_map>
#include <algorithm>
#include <climits>

/// @brief 最小堆中的元素，记录了某个编号的点此时的计数值
struct Elem
{
    int id = -1;// 点编号
    int cnt = 0;// 计数
    Elem() {}
    Elem(int _id, int _cnt): id(_id), cnt(_cnt) {}
    bool operator == (const Elem & b) const {
        return id == b.id;
    }
    bool operator < (const Elem & b) const {
        if (cnt == b.cnt) return id < b.id;
        else return cnt < b.cnt;
    }
};

/// @brief 最小堆
/// @tparam idx_t 索引类型
template<typename idx_t>
class MinHeap
{
private:
    bool assembled = false;// 是否已经完成建堆
    idx_t capa = 0;// 容量
    Elem * heapDataVec = nullptr;// 堆的数据存储向量
    idx_t * posit_table = nullptr;// 每个特定序号的元素当前存放在堆的什么位置
public:
    idx_t numCounts = 0;// 堆中元素个数
    MinHeap() {}
    MinHeap(const idx_t num) { initData(num, nullptr); }
    void printPos() {
        for (idx_t i = 0; i < capa; i++) {
            printf(" %3d ", posit_table[i]);
        }
        printf("\n");
    }
    /// @brief 向堆中插入新的元素
    /// @param elem 待插入元素
    void insert(Elem elem) {
        assert(numCounts < capa);
        assert(elem.id == numCounts);
        heapDataVec [numCounts] = elem;
        posit_table [elem.id] = numCounts;
        numCounts ++;
        percolateUp(numCounts - 1, heapDataVec[numCounts - 1]);// 对该元素做上滤
    }
    /// @brief 上滤
    /// @param index 元素的索引
    /// @param value 待上滤的元素
    void percolateUp(idx_t index, Elem value) {
        idx_t parentIndex = (index - 1) >> 1;
        while (index > 0 && value < heapDataVec[parentIndex]) {// 如果当前位置(由index指示)还不是堆顶，且父节点的值比本位置要大，则继续
            idx_t nameParent = heapDataVec[parentIndex].id;
            posit_table[nameParent] = index;// 更新原父节点的位置
            heapDataVec [index] = heapDataVec [parentIndex];
            index = parentIndex;
            parentIndex = (parentIndex - 1) >> 1;
        }
        // 找到了最终放置的位置
        heapDataVec[index] = value;
        posit_table[value.id] = index;
    }
    /// @brief 删除当前堆中的最小值（也就是堆顶元素）
    /// @return 返回被删除的元素
    Elem delMin() {
        Elem value = heapDataVec[numCounts - 1];// 原末尾元素
        Elem ret   = heapDataVec[0];// 堆顶元素，准备删除
        idx_t name_0 = ret.id; assert(posit_table[name_0] == 0);
        heapDataVec[numCounts - 1] = heapDataVec[0];// 原堆顶元素放到最后（即无效）
        posit_table[name_0] = numCounts - 1;// 记为

        numCounts --;// 堆中元素个数减一
        percolateDown(0, value);// 下滤，value是替换后的新的根节点值
        return ret;
    }
    /// @brief 下滤
    /// @param childTree 子节点的位置 
    /// @param value 待下滤的元素
    void percolateDown(idx_t childTree, Elem value) {
        // printf("percoDown id %d cnt %d\n", value.id, value.cnt);
        idx_t index = childTree;// 目标节点索引
        idx_t maxChid = (index + 1) << 1;// 目标节点右子节点
        bool godown = true;// 循环下滤退出标志
        while (godown && maxChid < numCounts) {// 右子节点尚还处于堆内
            godown = false;
            // printf("  maxChid %d     while [0] = id %d cnt %d\n", maxChid, heapDataVec[0].id, heapDataVec[0].cnt);
            if (heapDataVec[maxChid - 1] < heapDataVec[maxChid]) {// 若左孩子小，则更新为左孩子
                maxChid --;
                // printf("  maxChid %d     while [0] = id %d cnt %d\n", maxChid, heapDataVec[0].id, heapDataVec[0].cnt);
            }
            if (heapDataVec[maxChid] < value) {// 若父节点大于孩子，则下滤，此时孩子的位置记录需要更新
                godown = true;
                idx_t nameChid = heapDataVec[maxChid].id;
                posit_table[nameChid] = index;// 更新孩子的位置
                heapDataVec[index] = heapDataVec[maxChid];// 令较小值为交换值（将孩子上移）
                index = maxChid;// 该交换节点索引下移  
                maxChid = (maxChid + 1) << 1;// 重新计算交换节点右子节点  
            }
        }// while loop
        if (maxChid == numCounts) {// 右子节点已经不在堆内了
            if (heapDataVec[maxChid - 1] < value) {// 令左子节点值为交换值
                idx_t nameChid = heapDataVec[maxChid - 1].id;
                posit_table[nameChid] = index;// 更新孩子的位置
                heapDataVec[index] = heapDataVec[maxChid - 1];
                index = maxChid - 1;
            }
        }
        // printf("  final %d     while [0] = id %d cnt %d\n", index, heapDataVec[0].id, heapDataVec[0].cnt);
        // 找到了最终放置的位置
        heapDataVec[index] = value;// 将调整值赋予交换节点
        posit_table[value.id] = index;// 更新元素位置
    }
    /// @brief 建堆
    void makeHeap() {// 建堆的过程就是一个不断调整堆的过程，循环调用函数percolateDown调整子树  
        if (numCounts < 2) return;
        // 第一个需要调整的子树的根节点地址
        idx_t parent = (numCounts >> 1) - 1;
        while (true) {
            percolateDown(parent, heapDataVec[parent]);
            if (0 == parent) return; // 到达根节点，返回
            -- parent;
        }
    }
    /// @brief 查看堆顶元素
    /// @return 返回堆顶元素（的一份拷贝）
    Elem heapTop() {
        return heapDataVec[0]; 
    }
    /// @brief 初始化数据
    /// @param _num 数据个数
    /// @param _data 数据向量
    void initData(const idx_t _num, const Elem * _data) {
        assert(numCounts == 0 && capa == 0);
        numCounts = capa = _num;
        heapDataVec = new Elem [numCounts];
        posit_table = new idx_t[numCounts];
        if (_data == nullptr) {
            for (idx_t j = 0; j < numCounts; j++) {
                heapDataVec[j].id = j;
                posit_table[j] = j;
            }
        } else {
            for (idx_t j = 0; j < numCounts; j++) {
                heapDataVec[j] = _data[j];
                assert(_data[j].id < numCounts);
                posit_table[_data[j].id] = j;// 记录位于数据向量中的哪个位置
            }
        }
    }
    Elem * data() { return heapDataVec; }
    /// @brief 组装（包含建堆操作）
    void assemble() {
        assert(assembled == false);
        makeHeap();
        assembled = true;
    }
    /// @brief 将堆中某个元素的计数值（也就是权重）减一
    /// @param num 计数值要减一的元素个数
    /// @param names 它们的编号
    void decre(const idx_t num, const idx_t * names) {
        for (idx_t j = 0; j < num; j++) {
            const idx_t name = names[j];
            const idx_t posi = posit_table[name];
            // printf("decre %d at %d\n", name, posi);
            if (posi >= numCounts) continue;// 这个邻居之前已经被处理过了

            assert(heapDataVec[posi].id == name);
            heapDataVec[posi].cnt --;// 计数减一
            percolateUp(posi, heapDataVec[posi]);// 然后做上滤
        }
    }
    /// @brief  删除特定序号(名字)的元素，通过将其计数设为INT_MIN(最小的计数值)，然后上滤，则一定会出现在堆顶，再从堆顶删掉
    /// @param name 待删除的元素编号
    void delName(const idx_t name) {
        const idx_t posi = posit_table[name]; assert(posi < numCounts); assert(heapDataVec[posi].id == name);
        heapDataVec[posi].cnt = INT_MIN;
        percolateUp(posi, heapDataVec[posi]);
        Elem toDel = delMin(); assert(toDel.id == name);// 确认从堆顶删除的确实是需要删除的那个元素
    }
    /// @brief 析构函数
    ~MinHeap() {
        delete [] heapDataVec; heapDataVec = nullptr;
        delete [] posit_table; posit_table= nullptr;
        numCounts = capa = 0;
    }
};

/// @brief 基于点坐标的结对粗化方法
/// @tparam idx_t 索引类型
/// @tparam data_t 存储数据的浮点类型
/// @tparam calc_t 执行计算时的浮点类型
/// @param fine_mat 细网格矩阵
/// @param threshold 决定本点哪些邻居是强耦合/影响的阈值，一般为0.25（仅当使用基于最小堆的粗化方法时起效）
/// @param alpha 通过AC=R*AF*P计算粗网格矩阵AC时，一般设为1.0即可
/// @return 封装了粗化后的插值矩阵、限制矩阵、和粗网格矩阵的结构体
template<typename idx_t, typename data_t, typename calc_t> static
CoarsenWrapper<idx_t, data_t, calc_t>* pairwise_coarsen(const MyParCSRMatrix<idx_t, data_t, calc_t> & fine_mat,
    const calc_t threshold, const calc_t alpha)
{
    assert(threshold > 0.0);
    // int my_pid; MPI_Comm_rank(fine_mat.comm, &my_pid);
    // int nprocs; MPI_Comm_size(fine_mat.comm, &nprocs);
    // if (my_pid == 0) printf("Exact coarsen\n");

    const idx_t num_fine = fine_mat.end_row - fine_mat.beg_row;// 本进程在细网格所负责的行数
    MinHeap<idx_t> minheap(num_fine);// 最小堆
    Elem * ST_cnts = minheap.data();// 记录计数值的向量

    const MyCSRMatrix<idx_t, data_t, calc_t> & fine_diag = *(fine_mat.diag);// 细网格矩阵的对角部分
    const MyCSRMatrix<idx_t, data_t, calc_t> & fine_offd = *(fine_mat.offd);// 细网格矩阵的非对角部分
    idx_t * loc_map_f2c = new idx_t [num_fine];// 记录本地的第i个细点被映射为粗点的序号

    idx_t thread_flags [omp_get_max_threads()];// 多线程并行的标志位，用来表达该线程当前进行到哪个阶段
    idx_t thread_shared[omp_get_max_threads()];// 多线程相互共享的数据区
    for (idx_t t = 0; t < omp_get_max_threads(); t++) thread_flags[t] = -1;// 阶段编号初始化为-1
    idx_t * S_up_ptr = new idx_t [num_fine + 1], * S_upstrm = nullptr;// 记录细点的所有强耦合/影响的邻居个数和编号
    S_up_ptr[0] = 0;
    #pragma omp parallel
    {
        int tid = omp_get_thread_num();
        int nt  = omp_get_num_threads();
        int tbeg = 0, tend = 0;
        thread_load_balance(num_fine, nt, tid, tbeg, tend);// 将各线程的负载调整均衡

        idx_t * tl_rpt = new idx_t [tend - tbeg + 1];// 线程内部的行偏移
        tl_rpt[0] = 0;
        std::vector<idx_t> tl_S_upstrm;// 线程内部的强耦合/影响的邻居编号
        for (idx_t i = tbeg; i < tend; i++) {// 遍历本线程负责的每一行
            loc_map_f2c[i] = -1;
            const idx_t num_diag = fine_diag.row_ptr[i + 1] - fine_diag.row_ptr[i];
            // 决定某个邻居是否强影响我
            data_t min_neg = 0.0;
            // 检查对角部分的邻居
            for (idx_t p = fine_diag.row_ptr[i]; p < fine_diag.row_ptr[i + 1]; p++) {
                if (fine_diag.col_idx[p] == i) {
                    assert(fine_diag.vals[p] > 0.0);// 对角元要求为正
                    continue;// 对角元跳过
                }
                if (fine_diag.vals[p] < 0.0) {// 原算法好像只对负的非对角元处理？
                    min_neg = MIN(min_neg, fine_diag.vals[p]);
                }
            }
            // 检查非对角部分的邻居
            for (idx_t p = fine_offd.row_ptr[i]; p < fine_offd.row_ptr[i + 1]; p++) {
                if (fine_offd.vals[p] < 0.0) {// 原算法好像只对负的非对角元处理？
                    min_neg = MIN(min_neg, fine_offd.vals[p]);
                }
            }
            // 定义本点受其他点影响强弱的阈值（是一个负数）
            const data_t beta = threshold * min_neg; assert(beta < 0.0);
            data_t vals[num_diag];
            idx_t cnt = 0;
            // 根据阈值，确定强连通的邻居
            // 注意只需要在对角矩阵部分找强连通邻居即可，非对角部分 (non-local) 不用管
            for (idx_t p = fine_diag.row_ptr[i]; p < fine_diag.row_ptr[i + 1]; p++) {
                const idx_t ngb_id = fine_diag.col_idx[p]; assert(ngb_id < num_fine);
                if (ngb_id == i) continue;
                // if (nprocs == 1) if (ngb_id / (num_fine/2) != i / (num_fine/2)) continue;// 模拟多个进程的情况
                if (fine_diag.vals[p] < beta) {
                    tl_S_upstrm.push_back(ngb_id);// S_upstrm.push_back(ngb_id);
                    vals [cnt++] = fine_diag.vals[p];
                    __atomic_fetch_add( & ST_cnts[ngb_id].cnt, 1, __ATOMIC_RELAXED);// ST_cnts[ngb_id].cnt ++;
                }
            }
            const idx_t pend = tl_S_upstrm.size(),
                        pbeg = pend - cnt;
            idx_t * idxs = tl_S_upstrm.data() + pbeg;
            // 根据非零元的(负)值从小到大（即绝对值从大到小）冒泡排序 S_upstrm[pbeg, pend) ~ vals[0, cnt)
            bubble_sort<data_t, idx_t>(cnt, vals, idxs);

            // 记录这个点有多少个强影响的邻居
            tl_rpt[i - tbeg + 1] = pend;// S_up_ptr.push_back(S_upstrm.size());
        }
        assert(tl_rpt[tend - tbeg] == (idx_t) tl_S_upstrm.size());
        
        // 多线程数据汇总
        while (tid > 0 && __atomic_load_n( & thread_flags[tid - 1], __ATOMIC_ACQUIRE) < 1) {  }// 等待上一个线程完成
        idx_t prev_tsum = (tid > 0) ? thread_shared[tid - 1] : 0;
        thread_shared[tid] = prev_tsum + tl_rpt[tend - tbeg];// 本线程内的求和存入多线程共享区
        __atomic_store_n( & thread_flags[tid], 1, __ATOMIC_RELEASE);

        for (idx_t i = tbeg; i < tend; i++) {// 填充本线程负责范围内的行偏移
            S_up_ptr[i + 1] = prev_tsum + tl_rpt[i - tbeg + 1];
        }
        __atomic_store_n( & thread_flags[tid], 2, __ATOMIC_RELEASE);
        
        if (tid == 0) {
            while (__atomic_load_n( & thread_flags[nt - 1], __ATOMIC_ACQUIRE) < 2) {  }// 等待最后一个线程填充完
            const idx_t nnz = S_up_ptr[num_fine];// printf("nnz %d\n", nnz);
            S_upstrm = new idx_t [nnz];
            __atomic_store_n( & thread_flags[tid], 3, __ATOMIC_RELEASE);
        }
        
        while (tid > 0 && __atomic_load_n( & thread_flags[tid - 1], __ATOMIC_ACQUIRE) < 2) {  }// 等待上一个线程把假定的(assum)行偏移填充完
        while (tid > 0 && __atomic_load_n( & thread_flags[  0    ], __ATOMIC_ACQUIRE) < 3) {  }// 等待0号线程开辟好内存空间
        const idx_t dst_off = S_up_ptr[tbeg], num_write = tl_S_upstrm.size();
        for (idx_t j = 0; j < num_write; j++)
            S_upstrm[dst_off + j] = tl_S_upstrm[j];
    }// implicit sync

    idx_t num_coar = 0;// 本进程内生成的粗点数目
    idx_t * assum_R_rp = new idx_t [num_fine + 1];// 限制矩阵的行偏移，按最大可能空间开辟
    idx_t * assum_R_ci = new idx_t [num_fine];// 限制矩阵的列索引，它的非零元数一定是细网格的数目
    assum_R_rp[0] = 0;
    minheap.assemble();// 查询/修改前必须先完成组装（堆的建立）

    // 选粗点对
    while (minheap.numCounts > 0) {
        const idx_t id = minheap.heapTop().id;// printf("id %d cnt %d\n", id, minheap.heapTop().cnt);
        assert(loc_map_f2c[id] == -1);// 之前一定还未聚合过该点

        const idx_t pbeg = S_up_ptr[id], pend = S_up_ptr[id + 1];
        idx_t aggreg_id = -1;// 准备和哪个点进行聚合？
        idx_t offset = assum_R_rp[num_coar];
        loc_map_f2c[id] = num_coar;// 更新插值表
        assum_R_ci[offset ++] = id;// 更新限制矩阵
        minheap.delMin();
        minheap.decre(pend - pbeg, & S_upstrm[pbeg]);

        for (idx_t p = pbeg; p < pend; p++) {
            idx_t tenta_id = S_upstrm[p];// tentative aggregate
            if (loc_map_f2c[tenta_id] >= 0) {// 强影响本点的这个点之前已经被聚合过了
                continue;// 检查下一个
            }
            else {// 找到了尚未被聚合，且确实强影响本id点的
                aggreg_id = tenta_id;// 记录一下要聚合的点序号
                loc_map_f2c[aggreg_id] = num_coar;// 更新插值表
                assum_R_ci[offset ++] = aggreg_id;// 更新限制矩阵
                minheap.delName(aggreg_id);

                minheap.decre(S_up_ptr[aggreg_id + 1] - S_up_ptr[aggreg_id], & S_upstrm[S_up_ptr[aggreg_id]]);
                break;// 然后跳出 本点 的循环
            }
        }
        num_coar ++;// 粗点数增加
        assum_R_rp[num_coar] = offset;
    }
    assert(assum_R_rp[num_coar] == num_fine);
    assert(num_coar > 0);// 暂时不处理跨进程粗化
    delete [] S_up_ptr; S_up_ptr = nullptr;
    delete [] S_upstrm; S_upstrm = nullptr;

// #define REORDER_PR // 是否对生成的限制和插值矩阵进行重排
#ifdef REORDER_PR
    // 将细点编号小的往前挪，使其粗点编号也较小
    struct F2OldC {
        idx_t minF = -1;
        idx_t oldC = -1;
        bool operator < (const F2OldC & b) const {
            assert(minF != b.minF);
            return minF  < b.minF;
        }
    };
    std::vector<F2OldC> aux_arr(num_coar);
    for (idx_t oci = 0; oci < num_coar; oci++) {
        aux_arr[oci].oldC = oci;
        idx_t minF = INT_MAX;
        for (idx_t p = assum_R_rp[oci]; p < assum_R_rp[oci + 1]; p++)
            minF = MIN(minF, assum_R_ci[p]);
        aux_arr[oci].minF = minF;
    }
    std::sort(aux_arr.begin(), aux_arr.end());
    idx_t * rstr_rpt = new idx_t [num_coar + 1];
    idx_t * rstr_cid = new idx_t [num_fine];
    rstr_rpt[0] = 0;
    for (idx_t nci = 0; nci < num_coar; nci++) {
        idx_t oci = aux_arr[nci].oldC;
        for (idx_t p = assum_R_rp[oci]; p < assum_R_rp[oci + 1]; p++) {
            idx_t fi = assum_R_ci[p];
            loc_map_f2c[fi] = nci;// 重新构造细到粗的映射（相当于插值矩阵）
        }
        // 同时修改限制矩阵的行偏移
        rstr_rpt[nci + 1] = assum_R_rp[oci + 1] - assum_R_rp[oci];
    }
    for (idx_t nci = 0; nci < num_coar; nci++)
        rstr_rpt[nci + 1] += rstr_rpt[nci];
    // 修改限制矩阵的列索引
    for (idx_t nci = 0; nci < num_coar; nci++) {
        idx_t oci = aux_arr[nci].oldC;
        const idx_t num = assum_R_rp[oci + 1] - assum_R_rp[oci];
        idx_t * dst = rstr_cid + rstr_rpt[nci];
        const idx_t * src = assum_R_ci + assum_R_rp[oci];
        for (idx_t p = 0; p < num; p++)
            dst[p] = src[p];
    }
    delete [] assum_R_rp; assum_R_rp = nullptr;
    delete [] assum_R_ci; assum_R_ci = nullptr;
#else
    idx_t * rstr_rpt = assum_R_rp;// 限制矩阵的行偏移
    idx_t * rstr_cid = assum_R_ci;// 限制矩阵的列索引
#endif
    // 至此已经获得了聚合的列表
    // 可以计算粗网格的矩阵元素
    MyParCSRMatrix<idx_t, data_t, calc_t> * coar_mat = calc_coar_mat<idx_t, data_t, calc_t>
        (fine_mat, num_coar, rstr_rpt, rstr_cid, loc_map_f2c, alpha);

    // 返回的所有粗化信息
    CoarsenWrapper<idx_t, data_t, calc_t> * ret = new CoarsenWrapper<idx_t, data_t, calc_t>;
    ret->loc_map_f2c = loc_map_f2c;
    ret->R_rp = rstr_rpt;
    ret->R_ci = rstr_cid;
    ret->coar_mat = coar_mat;

    return ret;
}

#endif